/*
 ***************************************************************************
 * \file   convdiff1.cpp
 * \par Description
 * This code solves the set of equations:
 * \f[ 
 * \frac{\partial T}{\partial t} + 
 * u_j \frac{\partial T}{\partial x_j} = \nabla^2 T
 * \f]. 
 * \par
 * These equation are solved in a unit square \f$ x, y \in [0,1] \f$. 
 * The boundary conditions are as shown in the next figure:
 *
 * \image html convection.png "Unit square in 2D" width=10cm
 * \image latex convection.eps "Unit square in 2D" width=10cm
 * \par
 *  To compile and run this example type the next commands:
 * \par
   \verbatim
   % make
   % ./convdiff1 < input1 \endverbatim
 * 
 ***************************************************************************
 * \author    Luis M. de la Cruz [ Sat May 23 12:06:36 CDT 2009 ]
 ***************************************************************************
 */
// Note: Schemes: Quick_..., CDS_..., Upwind_...
#include <iostream>
#include <cmath>

#include "Meshes/Uniform.hpp"
#include "Equations/ScalarEquation.hpp"
#include "Schemes/CDS_CoDi.hpp"
#include "Equations/Momentum_XCoDi.hpp"
#include "Schemes/CDS_XCoDi.hpp"
#include "Equations/Momentum_YCoDi.hpp"
#include "Schemes/CDS_YCoDi.hpp"
#include "Equations/PressureCorrection.hpp"
#include "Schemes/Simplec.hpp"
#include "Solvers/TDMA.hpp"
#include "Utils/inout.hpp" 
#include "Utils/num_utils.hpp"
#include "Utils/GNUplot.hpp"
#include "DDM/CartComm.hpp"
#include "DDM/SubDomain.hpp"
using namespace Tuna;


int FLAG;

// Scalar field in one dimension
typedef TunaArray<double,2>::huge ScalarField2D;

// Parameters to make the partition
double length_x, length_y;
int num_nodes_x, num_vols_x, num_subdom_x;
int num_nodes_y, num_vols_y, num_subdom_y;
int nc_ovlp_l, nc_ovlp_r, nc_ovlp_d, nc_ovlp_u;
double ovlp_r, ovlp_l, ovlp_u, ovlp_d;
double shiftx, shifty;

// Parameters to identify the subdomain
int Isub, Jsub, rank, size; 
int ddm_iter;   

// Parameters for the problem
double leftwall, rightwall;
double tolerance, dt;    
int tdma_iter, steps, frequency;
double Prandtl, Rayleigh;

double errorT, errorX, errorY;

bool read_data_and_Bcast(CartComm<2>&);

template<class T_e, class T_x, class T_y, class T_p>
double SIMPLEC(T_e &, T_x &, T_y &, T_p &, int, double, double, double,
	       CartComm<2>&);

int main( int argc, char **argv)
{

  //    MPI::Init(argc, argv);
  //    rank = MPI::COMM_WORLD.Get_rank();       
  //    size = MPI::COMM_WORLD.Get_size(); 
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
// ----- Construction of a Cartesian topology

    CartComm<2> cart(argc, argv, rank);   
    Isub = cart.get_I();
    Jsub = cart.get_J();
    num_subdom_x = cart.getNumProc_I();
    num_subdom_y = cart.getNumProc_J(); 
    //    cart.print();

// Reading and broadcasting data from the input file 
    read_data_and_Bcast(cart); 

// Global Mesh 
    StructuredMesh<Uniform<double, 2> > mesh(length_x, num_nodes_x,
					     length_y, num_nodes_y);    
    double shiftx = Isub * length_x / num_subdom_x;
    double shifty = Jsub * length_y / num_subdom_y;

// Subdomain
    SubDomain<double, 2> subdom(cart, mesh, 
				nc_ovlp_l, nc_ovlp_r,
				nc_ovlp_d, nc_ovlp_u);

    double dx = mesh.getDelta(X);
    double dy = mesh.getDelta(Y);
    int num_nodes_x = mesh.getExtentNodes(X);
    int num_nodes_y = mesh.getExtentNodes(Y);
    int num_vols_x = mesh.getExtentVolumes(X);
    int num_vols_y = mesh.getExtentVolumes(Y);

    //    mesh.print(Isub, Jsub);
    //subdom.print();


    ScalarField2D T(mesh.getExtentVolumes());      // Temperature
    ScalarField2D p(mesh.getExtentVolumes());      // pressure
    ScalarField2D pp(mesh.getExtentVolumes());     // pressure correction
    ScalarField2D u(mesh.getExtentVolumes());      // u-velocity
    ScalarField2D v(mesh.getExtentVolumes());      // v-velocity
    ScalarField2D us(num_nodes_x, num_vols_y); // u-velocity staggered
    ScalarField2D vs(num_vols_x, num_nodes_y); // v-velocity staggered

    ScalarField2D Tn(mesh.getExtentNodes());
    ScalarField2D pn(mesh.getExtentNodes());
    ScalarField2D un(mesh.getExtentNodes());
    ScalarField2D vn(mesh.getExtentNodes());    

    Range all = Range::all();

    T = 0; Tn = 0;
    p = 0; pn = 0;
    u = 0; un = 0;
    v = 0; vn = 0;

    if (Isub != 0) shiftx -= dx * nc_ovlp_l;
    if (Jsub != 0) shifty -= dy * nc_ovlp_d; 
    if (Isub == 0)
      T(T.lbound(firstDim), all) = leftwall;   // Left
    if (Isub == num_subdom_x - 1)
      T(T.ubound(firstDim), all) = rightwall;    // Rigth    

    SparseMatrix< Diagonal<double, 2> > A(num_nodes_x, num_nodes_y); 
    ScalarField2D                       b(num_nodes_x, num_nodes_y);


//
// Energy equation definition
//
    ScalarEquation<CDS_CoDi<double,2> > energy(T, A, b, mesh.getDeltas());
    energy.setDeltaTime(dt);
    energy.setUvelocity(us);
    energy.setVvelocity(vs);
    subdom.setRealDirichlet(energy, LEFT_WALL);
    subdom.setRealDirichlet(energy, RIGHT_WALL);
    subdom.setRealNeumann(energy, TOP_WALL);
    subdom.setRealNeumann(energy, BOTTOM_WALL);
    subdom.setGhostDirichlet(energy);
    //    energy.print(rank);

//
// Xmomentum equation definition
//
    Momentum_XCoDi<CDS_XCoDi<double, 2> > mom_x(us, A, b, mesh.getDeltas());
    mom_x.setGamma(Prandtl);
    mom_x.setDeltaTime(dt);
    subdom.setRealDirichlet(mom_x, LEFT_WALL);
    subdom.setRealDirichlet(mom_x, RIGHT_WALL);
    subdom.setRealDirichlet(mom_x, TOP_WALL);
    subdom.setRealDirichlet(mom_x, BOTTOM_WALL);
    subdom.setGhostDirichlet(mom_x);
    mom_x.setUvelocity(us);
    mom_x.setVvelocity(vs);
    mom_x.setPressure(p);
    //    mom_x.print(rank);

//
// Ymomentum equation definition
//    
    Momentum_YCoDi<CDS_YCoDi<double, 2> > mom_y(vs, A, b, mesh.getDeltas());
    mom_y.setGamma(Prandtl);
    mom_y.setDeltaTime(dt);
    subdom.setRealDirichlet(mom_y, LEFT_WALL);
    subdom.setRealDirichlet(mom_y, RIGHT_WALL);
    subdom.setRealDirichlet(mom_y, TOP_WALL);
    subdom.setRealDirichlet(mom_y, BOTTOM_WALL);
    subdom.setGhostDirichlet(mom_y);
    mom_y.setUvelocity(us);
    mom_y.setVvelocity(vs);
    mom_y.setPressure(p);
    mom_y.setTemperature(T);
    mom_y.setRayleigh(Rayleigh);
    //    mom_y.print(rank);


//
// Pressure correction equation definition
// 
    PressureCorrection<Simplec<double, 2> > press(pp, A, b, mesh.getDeltas());
    press.applyBounds(1, num_nodes_x-1, 1, num_nodes_y-1);
    press.setUvelocity(us);
    press.setVvelocity(vs);
    press.set_du(mom_x.get_du());
    press.set_dv(mom_y.get_dv());
    press.setPressure(p);
    //    press.print(rank);

    /***
    ofstream sorsumfile("./Data1/sorsum"), nusseltfile("./Data1/nusselt"),
      errorTfile("./Data1/errorT"), errorXfile("./Data1/errorX"), 
      errorYfile("./Data1/errorY");
    ***/
    double sorsum, resorm, nusselt, global_sorsum = 1000;
    
    int iteration = 0;

    NumUtils::interpolateToNodes(Tn, T);
    NumUtils::interpolateToNodes(pn, p);
    NumUtils::interpolateToNodesUs(un, us);
    NumUtils::interpolateToNodesVs(vn, vs);

    ostringstream dir_number;
    dir_number.width(2);
    dir_number.fill('0');
    dir_number << rank;
    string dnumber = dir_number.str();
    string filenameV = "./DataConvDiff2D/" + dnumber + "/velc.";
    string filenameT = "./DataConvDiff2D/" + dnumber + "/temp.";
    string filenameP = "./DataConvDiff2D/" + dnumber + "/pres.";

    InOut::writeToFile_DXP(un, vn, 0, filenameV, dx, shiftx, dy, shifty);
    InOut::writeToFile_DXP(Tn, 0, filenameT, dx, shiftx, dy, shifty);
    InOut::writeToFile_DXP(pn, 0, filenameP, dx, shiftx, dy, shifty);

    /*** ***/

    //    while ( (global_sorsum > tolerance) && (iteration < steps) ) {

    for(iteration = 1; iteration <= steps; ++iteration) {

      sorsum = SIMPLEC(energy, mom_x, mom_y, press,
		       tdma_iter, tolerance, dx, dy, cart);

      // Exchange of info
      subdom.infoExchange(T);
      energy.updatePhi(T);

      NumUtils::interpolateToNodesUs(un, us);
      NumUtils::interpolateToNodesVs(vn, vs);
      subdom.infoExchange(un, 0);
      subdom.infoExchange(vn, 0);
      NumUtils::staggerFromNodesX(us, un);
      NumUtils::staggerFromNodesY(vs, vn);
      mom_x.updatePhi(us);
      mom_y.updatePhi(vs);

      subdom.infoExchange(p);
      press.updatePhi(p); 
      
      if ( (iteration % frequency) == 0 ) {
	NumUtils::interpolateToNodes(Tn, T);
	NumUtils::interpolateToNodes(pn, p);
	//	NumUtils::interpolateToNodesUs(un, us);
	//	NumUtils::interpolateToNodesVs(vn, vs);	    
	InOut::writeToFile_DXP(un, vn, iteration, filenameV, dx, shiftx, dy, shifty);
	InOut::writeToFile_DXP(Tn, iteration, filenameT, dx, shiftx, dy, shifty);
	InOut::writeToFile_DXP(pn, iteration, filenameP, dx, shiftx, dy, shifty);
      }

  
      if (rank == FLAG)
	cout << "\n => Iter = " << iteration << " sorsum = " << sorsum;
      	
      /*** *
      sorsumfile << iteration << "\t" << sorsum << endl;     
      nusselt = NumUtils::calcNusseltVert(T,  dx * 0.5, dy, 0);
      nusseltfile << iteration << "\t" << nusselt << endl;
      errorTfile << iteration << "\t" << errorT << endl;
      errorXfile << iteration << "\t" << errorX << endl;
      errorYfile << iteration << "\t" << errorY << endl;
      /* ***/

 
    }

    /*** ***/

    MPI_Finalize();


    return 0;
}


template<class T_e, class T_x, class T_y, class T_p>
double SIMPLEC(T_e &energy, T_x &mom_x, T_y &mom_y, T_p &pressure, 
	       int max_iter, double tolerance, double dx, double dy,
	       CartComm<2>& cart) 
{
  double sorsum = 10.0, tol_simplec = 1e-2, local_sorsum;
    int counter = 0;

    while ( (sorsum > tol_simplec) && (counter < 1) ) {
	energy.calcCoefficients();
	Solver::TDMA2DY(energy, tolerance, max_iter);
	errorT = energy.calcErrorL1();
	energy.update();

	/*** ***/
	mom_x.calcCoefficients();	
	Solver::TDMA2DY(mom_x, tolerance, max_iter);
	errorX = mom_x.calcErrorL1();
	mom_x.update();

	mom_y.calcCoefficients();
	Solver::TDMA2DY(mom_y, tolerance, max_iter);
	errorY = mom_y.calcErrorL1();
	mom_y.update();

	pressure.calcCoefficients();
	Solver::TDMA2DY(pressure, tolerance, max_iter);
	pressure.correction();
	local_sorsum = fabs( pressure.calcSorsum() );

	//	MPI::COMM_WORLD.Allreduce(&local_sorsum, &sorsum, 
	//				  1, MPI::DOUBLE, 
	//				  MPI::MAX);


	/* ***/
	if (rank == FLAG) {
	  cout << "\n Cont = " << counter << "\t rank = " << rank << "\t eT = " << errorT 
	       << "\t eX = " << errorX << "\t eY = " << errorY 
	       << "\t local sorsum = " << local_sorsum;
	}
	/*** */

	MPI_Allreduce(&local_sorsum, &sorsum, 1, MPI_DOUBLE, MPI_MAX, cart.comm);

	++counter;
    }

    return sorsum;
    
}




bool read_data_and_Bcast(CartComm<2>& cart) 
{
  int    data_Int[10];
  double data_Flt[8];

  if (rank == 0) {
    std::cout << "\n\n"
	      << " +----------------------------------------+\n"
	      << " |    TUNA FOR PDE SOLVING IN PARALLEL    |\n"
	      << " +----------------------------------------+\n"
	      << " | MASTER PROCESS rank = " << rank << "\n"
	      << " | No subdomains in x-axis = " << num_subdom_x << "\n"
	      << " | No subdomains in y-axis = " << num_subdom_y << "\n"
	      << " +----------------------------------------+\n"
	      << "\n";
	
// ----- Reading data from "input" file	
	std::ifstream input_cin ("inputConvDiffPar2D");
	input_cin >> data_Flt[0]   // length in x-axis
		  >> data_Flt[1]   // length in y-axis
		  >> data_Int[0]   // Num of nodes in x-axis for each subdomain
		  >> data_Int[1]   // Num of nodes in y-axis for each subdomain
		  >> data_Int[2]  // Overlapping left
		  >> data_Int[3]  // Overlapping right
		  >> data_Int[4]  // Overlapping down
		  >> data_Int[5]  // Overlapping up
		  >> data_Flt[2] // Left wall condition
		  >> data_Flt[3] // Right wall condition
		  >> data_Flt[4]   // Time step
		  >> data_Int[6]   // Max number of iterations
		  >> data_Flt[5]   // Tolerance
		  >> data_Int[7]   // Max iterations for TDMA algo
		  >> data_Flt[6]   // Prandtl
		  >> data_Flt[7]   // Rayleigh
		  >> data_Int[8]  // Print frequency
		  >> data_Int[9];  // FLAG
	input_cin.close();
    }
// ----- Broadcast the info to all processors
//    cart.comm.Bcast(data_Int, 10, MPI::INT, 0);    
//    cart.comm.Bcast(data_Flt, 5, MPI::DOUBLE, 0); 
  MPI_Bcast(data_Int, 10, MPI::INT, 0, cart.comm);    
  MPI_Bcast(data_Flt, 8, MPI::DOUBLE, 0, cart.comm);

// ----- Using global variables.
    length_x    = data_Flt[0];
    length_y    = data_Flt[1];
    leftwall    = data_Flt[2];
    rightwall   = data_Flt[3]; 
    dt          = data_Flt[4];
    tolerance   = data_Flt[5];
    Prandtl     = data_Flt[6];
    Rayleigh    = data_Flt[7];

    num_nodes_x = data_Int[0];
    num_nodes_y = data_Int[1];
    nc_ovlp_l   = data_Int[2]; 
    nc_ovlp_r   = data_Int[3];
    nc_ovlp_d   = data_Int[4]; 
    nc_ovlp_u   = data_Int[5];    
    steps       = data_Int[6];
    tdma_iter   = data_Int[7];
    frequency   = data_Int[8];
    FLAG        = data_Int[9];

    return 0;
}

